#加载预训练的 FCN 模型，将其转换为 Torch 脚本模块，并保存以备后续使用
import torch
# from UNet import UNet
from fcn import FCN
#设置所用设备为 CPU
device = torch.device("cpu")
#创建一个 FCN 类的实例，并设置模型输出的类别数为 1
seg_net = FCN(num_classes=1)
#加载预训练模型，这里使用的是第 20 个 epoch 的模型，并将加载的模型参数映射到 CPU 设备
seg_net.load_state_dict(torch.load("results/pth/198.pth", map_location=device))
#将模型设置为评估模式
seg_net.eval()
#创建一个随机张量用于 Torch 脚本模块的示例输入，代表了模型预期接收的输入格式(批处理大小，通道数，高度，宽度)
x = torch.rand(1, 3, 256, 256)
#将模型转换为 Torch 脚本模块
traced_script_module = torch.jit.trace(func=seg_net, example_inputs=x)
#保存 Torch 脚本模块
traced_script_module.save("198_fcn.pt")